package spinal_loongarch_core132

import spinal.core._
import spinal.lib._

trait CacheConfig{
    val w_idx :Int
    val w_pos :Int
    val n_way :Int
    val withDirty:Boolean
    val withBitMask:Boolean
    def n_bank:Int = 1 << w_pos
    def n_line:Int = 1 << w_idx
    def w_tagv:Int = 21
    def w_data:Int = 32+withDirty.toInt
    def getTagV = for(i <- 0 until n_way) yield new TagVBlockRam
    def getIDat = Range(0,n_way).map((i:Int)=>Range(0,n_bank).map((j:Int)=> new InstBlockRam))
    def getDDat = Range(0,n_way).map((i:Int)=>Range(0,n_bank).map((j:Int)=> new DataBlockRam))
    def each_way(op:Int=>Unit){
        for(i <- 0 until n_way){
            op(i)
        }
    }
    def each_blk(op:(Int,Int)=>Unit){
        for(i <- 0 until n_way){
            for(j <- 0 until n_bank){
                op(i,j)
            }
        }
    }
    def getWayBits(get:Int=>Bool):Bits = {
        Vec(Range(0,n_way).map((i:Int)=>
            get(i)
        )).asBits
    }
    def getVec1D(get:(Int)=>UInt) = {
        Vec(Range(0,n_way).map((i:Int)=>
            get(i)
        ))
    }
    def getVec2D(get:(Int,Int)=>UInt) = {
        Vec(Range(0,n_way).map((i:Int)=>
            Vec(Range(0,n_bank).map((j:Int)=>
                get(i,j)
            ))
        ))
    }
    def getWayMask(way:UInt):Bits = {
        return B(1, n_way bits) |<< way
    }
}

abstract class LookupInfo extends Bundle{
    val cfg:CacheConfig
    val pt   :UInt
    val mat  :UInt
    val ex   :Bool
    val ecode:UInt
    val vidx :UInt
    val cached:Bool 
    def ptag  :UInt = pt
    def tagv  :UInt = U("1") @@ ptag
    def widx  :UInt = vidx(4, 8 bits)
    def wpos  :UInt = vidx(2, 2 bits)
    def pa    :UInt = ptag @@ vidx
    val va    :UInt = LISA.GPR
    def ridx  :UInt = va(4, 8 bits)
    def rpos  :UInt = va(2, 2 bits)
    def nextVIdx:UInt = va(0, 12 bits)
    def tagCompare(tagv:IndexedSeq[TagVBlockRam]):(Bits,Bool) = {
        val hits:IndexedSeq[Bool] = tagv.map(_.io.rd === U("1") @@ ptag)
        val hit :Bool = hits.reduce(_||_)
        return (Vec(hits).asBits,hit)
    }
    def select(hits:Bits,data:Vec[Vec[UInt]]):UInt = {
        MuxOH0(hits.asBools,Range(0,cfg.n_way).map((i:Int)=>data(i)(wpos))).asUInt
    }
    def selectWay(hits:Bits,data:Vec[Vec[UInt]]):Vec[UInt] = {
        Vec(Range(0,cfg.n_bank).map((j:Int)=>
            MuxOH0(hits.asBools,Range(0,cfg.n_way).map((i:Int)=>data(i)(j))).asUInt
        ))
    }
    def send_ex(e:IExcept){
        e.ecode := ecode
    }
}
class LRUReplace extends Area{
    val ram = Reg(UInt(256 bits))
    val bit = Reg(Bool)
    val wa  = UInt(8 bits)
    val we  = Bool()
    val wd  = Bool()
    we := False
    wa := U(0)
    wd := False
    when(we){
        ram(wa) := wd
    }
    def read(idx:UInt):Unit = {
        bit := ram(idx)
    }
    def write(idx:UInt,bit:Bool):Unit = {
        wa := idx
        wd := bit
        we := True
    }
    def update(fill:Bool,hits:Bits):Bool = {
        return Mux(fill,!bit,hits(0))
    }
    def way:UInt = {
        return bit.asUInt
    }
}
class BusReceiver(resp:CacheResp,miss:IMiss) extends Bundle{
    object BusState extends SpinalEnum{
        val sIdle,sWaitData,sResponse = newElement()
    }
    import BusState._
    val s = Reg(BusState()) init sIdle
    val go    = Bool()
    val s_next = BusState()
    val data  = Reg(UInt(32 bits))
    val start = Bool() // from controller
    val valid = Bool()
    val done  = Bool()
    val cancel= resp.cancel

    // sIdle -- start --> sWaitData
    // -- receive the data --> sResponse
    // -- core takes this data --> sIdle
    valid := s === sResponse
    go := False
    switch(s){
        is(sIdle){
            done := True
            s_next := sWaitData
            when(start){
                go := True
            }
        }
        is(sWaitData){
            done := False
            s_next := sResponse
            when(miss.rvalid){
                go := True
                data   := miss.rdata
            }
        }
        is(sResponse){
            done := resp.allow
            s_next := sIdle
            when(resp.allow){
                go := True
            }
        }
    }
    when(go){
        s := s_next
    }
    if(cancel != null){
        when(cancel){
            s := sIdle
        }
    }
}
class IExcept extends Bundle{
    val ex    :Bool = Bool()
    val ecode :UInt = LISA.Ex.Code
}
class DExcept extends Bundle{
    val ex    :Bool = Bool()
    val ecode :UInt = LISA.Ex.Code
    val esubc :UInt = LISA.Ex.Subc
}
class CacheResp extends Bundle with IMasterSlave{
    val valid :Bool = Bool()
    val allow :Bool = Bool()
    val data  :UInt = LISA.GPR
    val cancel:Bool = Bool()
    override def asMaster():Unit = {
        out(valid,data)
        in (allow,cancel)
    }
}
class CacheReq extends Bundle with IMasterSlave{
    val valid  :Bool = Bool()
    val allow  :Bool = Bool()
    val va     :UInt = LISA.GPR
    override def asMaster():Unit = {
        out(valid,va)
        in(allow)
    }
}
class CacheFiller(miss:IMiss) extends Area{
    val valid:Bool = RegInit(False)
    val start:Bool = Bool()
    val done :Bool = !valid
    val num  = Reg(UInt(2 bits))
    when(start){
        valid := True
        num := U(0)
    }.elsewhen(valid && miss.rvalid){
        valid := !miss.rlast
        num := num + 1
    }
    def isWriteTagV(miss:IMiss):Bool = valid && miss.rvalid && miss.rlast
    def isWriteData(miss:IMiss):Bool = valid && miss.rvalid
    def getPos(info:LookupInfo):UInt = info.wpos + num
}

abstract class CacheCtrl extends Area{
    object CacheState extends SpinalEnum{
        val sIdle,sLookup,sMiss,sOperate = newElement()
    }
    import CacheState._
    // input fields
    val req_valid  = Bool()
    val miss_ready = Bool()
    val resp_ready = Bool()
    val recv_done  = Bool()
    val fill_done  = Bool()
    val cancel     = Bool()
    val cached     = Bool()
    val ex         = Bool()
    val hit        = Bool()

    val need_resp:Bool
    val need_fill:Bool

    val init = RamInit(8)
    val s  = RegInit(sIdle)

    val addr_ok = Bool()
    val find_ok = Bool()
    val data_ok:Bool
    val miss_valid = Bool()
    val send_miss  = miss_valid && miss_ready

    def acceptReq:Unit
    def sendMiss:Unit
    def quitable:Bool = s =/= sMiss

    addr_ok := False
    find_ok := False
    miss_valid := False
    switch(s){
        is(sIdle){
            acceptReq
        }
        is(sLookup){
            when(hit && cached || ex){
                find_ok := True
                when(resp_ready){
                    acceptReq
                }
            }.otherwise{
                sendMiss
            }
        }
        is(sMiss){
            when(fill_done && recv_done){
                acceptReq
            }
        }
    }

    def collect(req:CacheReq):Unit = {
        req_valid := req.valid
        req.allow := addr_ok
    }
    def collect(resp:CacheResp):Unit = {
        resp_ready := resp.allow
        cancel := resp.cancel
    }
    def collect(recv:BusReceiver):Unit = {
        recv_done := recv.done
        recv.start := send_miss && need_resp
    }
    def collect(fill:CacheFiller):Unit = {
        fill_done := fill.done
        fill.start := send_miss && need_fill
    }

}
